{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/agents/gemini_data_analytics/intro_gemini_data_analytics_http.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "cfNFHkaMKA97",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1762978722384,
          "user_tz": 300,
          "elapsed": 2,
          "user": {
            "displayName": "",
            "userId": ""
          }
        }
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jD4y5MNU_qi4"
      },
      "source": [
        "# Intro to Gemini Data Analytics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RKVaW_pICJ2R"
      },
      "source": [
        "| Author |\n",
        "| --- |\n",
        "| [Aditya Verma](https://github.com/vermaAstra) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3Z7ymg4eGyeU"
      },
      "source": [
        "# Background and Overview\n",
        "The **Conversational Analytics API** lets you chat with your BigQuery or Looker data anywhere, including embedded Looker dashboards, Slack and other chat apps, or even your own web applications. Your team members can get answers where they need them, when they need them, in the applications they use every day. You can find the [Colab example with the Python SDK here](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/gemini_data_analytics/intro_gemini_data_analytics_sdk.ipynb). This is a **Pre-GA** product. See [documentation](https://cloud.google.com/gemini/docs/conversational-analytics-api/overview) for more details.\n",
        "\n",
        "Please provide feedback to conversational-analytics-api-feedback@google.com\n",
        "<br>\n",
        "### This notebook will help you\n",
        "1. Authenticate to Google Cloud\n",
        "2. Add data\n",
        "3. Perform agent operations (create, list, get, delete)\n",
        "4. Manage conversations (create, list, get, delete)\n",
        "5. Ask questions with your agent\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d1divKxHBVyQ"
      },
      "source": [
        "# Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YGT_vK5nSrCB"
      },
      "source": [
        "## API Enablement\n",
        "\n",
        "**Please fill in the billing_project form field with your own Google Cloud project.  The project must have the following APIs enabled:**\n",
        "-  [cloudaicompanion API](https://console.cloud.google.com/apis/library/cloudaicompanion.googleapis.com)\n",
        "-  [Gemini Data Analytics API](https://console.cloud.google.com/apis/library/geminidataanalytics.googleapis.com)\n",
        "-  [BQ API](https://console.cloud.google.com/marketplace/product/google/bigquery.googleapis.com)\n",
        "-  [Dataform API](https://console.cloud.google.com/apis/library/dataform.googleapis.com)\n",
        "- [Vertex AI API](https://console.cloud.google.com/apis/library/aiplatform.googleapis.com)\n",
        "\n",
        "You may pass in any BigQuery project/dataset/table for which you have read permissions.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DZz5XBuNSRb0"
      },
      "source": [
        "## Auth and Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Skzx_KNaG3c-"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import json as json_lib\n",
        "\n",
        "import altair as alt\n",
        "import pandas as pd\n",
        "import requests\n",
        "from IPython.display import HTML, display\n",
        "from google.colab import auth\n",
        "from pygments import formatters, highlight, lexers\n",
        "import textwrap\n",
        "\n",
        "\n",
        "auth.authenticate_user()\n",
        "\n",
        "access_token = !gcloud auth application-default print-access-token\n",
        "headers = {\n",
        "    \"Authorization\": f\"Bearer {access_token[0]}\",\n",
        "    \"Content-Type\": \"application/json\",\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5nEg33LQ4AmH"
      },
      "source": [
        "# Datasource and Billing project\n",
        "\n",
        "\n",
        "*   Define your billing project within Google Cloud\n",
        "*   Link your agent to data sources (BigQuery, Looker, Looker Studio)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a1fb79c902a9"
      },
      "outputs": [],
      "source": [
        "# @title Billing Project and Prompt (System Instruction)\n",
        "\n",
        "# fmt: off\n",
        "billing_project = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "location = \"global\"\n",
        "api_version = \"v1beta\"\n",
        "\n",
        "# provide critical context for your Conversational Analytics Agent here\n",
        "system_instruction = \"Think like an Analyst\"  # @param {type:\"string\"}\n",
        "# fmt: on"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c46d107e7768"
      },
      "outputs": [],
      "source": [
        "# @title BigQuery Datasource\n",
        "\n",
        "bigquery_data_sources = {\n",
        "    \"bq\": {\n",
        "        \"tableReferences\": [\n",
        "            {\n",
        "                \"projectId\": \"bigquery-public-data\",\n",
        "                \"datasetId\": \"faa\",\n",
        "                \"tableId\": \"us_airports\",\n",
        "            },\n",
        "            {\n",
        "                \"projectId\": \"bigquery-public-data\",\n",
        "                \"datasetId\": \"san_francisco\",\n",
        "                \"tableId\": \"street_trees\",\n",
        "            },\n",
        "            # Add more table references here\n",
        "        ]\n",
        "    }\n",
        "}\n",
        "\n",
        "# optional example queries (only leveraged for BigQuery datasources currently)\n",
        "example_queries = [\n",
        "    {\n",
        "        \"naturalLanguageQuestion\": \"What is the highest observed positive longitude?\",\n",
        "        \"sqlQuery\": \"SELECT MAX(longitude) FROM airports\",\n",
        "    },\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ae7bc48aa40"
      },
      "outputs": [],
      "source": [
        "# @title Looker Datasource\n",
        "\n",
        "# When connecting to the Looker datasource, you can authenticate using either:\n",
        "# 1. client_id and client_secret\n",
        "# 2. access_token\n",
        "looker_credentials = {\n",
        "    \"oauth\": {\n",
        "        \"secret\": {\n",
        "            \"client_id\": \"Enter your Client ID Here\",\n",
        "            \"client_secret\": \"Enter your Client Secret Here\",\n",
        "        }\n",
        "    }\n",
        "}\n",
        "\n",
        "# Uncomment this and comment out the above one if you are using access_token for authentication\n",
        "# looker_credentials = {\n",
        "#     \"oauth\": {\n",
        "#         \"token\": {\n",
        "#           \"access_token\": \"your_looker_access_token\",\n",
        "#         }\n",
        "#     }\n",
        "# }\n",
        "\n",
        "# Looker datasource 1\n",
        "# fmt: off\n",
        "lookml_model_1 = \"<add lookml_model here>\"  # @param {type:\"string\"}\n",
        "explore_1 = \"<add explore here>\"  # @param {type:\"string\"}\n",
        "\n",
        "# Looker datasource 2 (optional)\n",
        "# lookml_model_2 = \"<add lookml_model here>\"  # @param {type:\"string\"}\n",
        "# explore_2 = \"<add explore here>\"  # @param {type:\"string\"}\n",
        "\n",
        "# Required only for Looker PUBLIC instance\n",
        "looker_instance_uri = \"https://my_company.looker.com\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "# Looker datasources for PUBLIC Instance\n",
        "looker_data_sources = {\n",
        "    \"looker\": {\n",
        "        \"explore_references\": [\n",
        "            {\n",
        "                \"looker_instance_uri\": looker_instance_uri,\n",
        "                \"lookml_model\": lookml_model_1,\n",
        "                \"explore\": explore_1,\n",
        "            },\n",
        "            # {\n",
        "            #     \"looker_instance_uri\": looker_instance_uri,\n",
        "            #     \"lookml_model\": lookml_model_2,\n",
        "            #     \"explore\": explore_2,\n",
        "            # },\n",
        "            # Add up to 5 total explore references here\n",
        "        ],\n",
        "        # NOTE -\n",
        "        # 1. Do not pass looker credentials during create data agent.\n",
        "        # 2. Uncomment the below line only when running the method [Stateless Chat using inline context]\n",
        "        # \"credentials\": looker_credentials\n",
        "    }\n",
        "}\n",
        "\n",
        "\n",
        "# # Looker datasource for PRIVATE Instance\n",
        "\n",
        "# looker_data_sources = {\n",
        "#     \"looker\": {\n",
        "#       \"explore_references\": [\n",
        "#           {\n",
        "#             \"private_looker_instance_info\": {\n",
        "#                 \"looker_instance_id\": \"your_private_looker_instance_id\",\n",
        "#                 \"service_directory_name\": \"your_private_looker_instance_service_directory_name\",\n",
        "#             },\n",
        "#             \"lookml_model\": lookml_model_1,\n",
        "#             \"explore\": explore_1,\n",
        "#           },\n",
        "#           # {\n",
        "#           #   \"private_looker_instance_info\": {\n",
        "#           #       \"looker_instance_id\": \"your_private_looker_instance_id\",\n",
        "#           #       \"service_directory_name\": \"your_private_looker_instance_service_directory_name\",\n",
        "#           #   },\n",
        "#           #   \"lookml_model\": lookml_model_2,\n",
        "#           #   \"explore\": explore_2,\n",
        "#           # },\n",
        "# .          # Add up to 5 total explore references here\n",
        "#       ],\n",
        "#       # NOTE -\n",
        "#       # 1. Do not pass looker credentials during create data agent.\n",
        "#       # 2. Uncomment the lines below only when running the method [Stateless Chat using inline context]\n",
        "#       # \"credentials\": {\n",
        "#       #     \"oauth\": {\n",
        "#       #         \"token\": {\n",
        "#       #           \"access_token\": \"your_looker_access_token\",\n",
        "#       #         }\n",
        "#       #     }\n",
        "#       # }\n",
        "#     }\n",
        "# }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "61cf3fddad05"
      },
      "outputs": [],
      "source": [
        "# @title Looker Studio Datasource\n",
        "\n",
        "looker_studio_data_sources = {\n",
        "    \"studio\": {\"studio_references\": [{\"datasource_id\": \"your_studio_datasource_id\"}]}\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "66e0d8583aa7"
      },
      "source": [
        "# Environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OE6bPKsyw_Jy"
      },
      "outputs": [],
      "source": [
        "# ENVIRONMENT OPTIONS: prod, autopush, staging\n",
        "# fmt: off\n",
        "environment = \"prod\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "base_url = \"https://geminidataanalytics.googleapis.com\"\n",
        "\n",
        "if environment == \"autopush\":\n",
        "    base_url = \"https://autopush-geminidataanalytics.sandbox.googleapis.com\"\n",
        "elif environment == \"staging\":\n",
        "    base_url = \"https://staging-geminidataanalytics.sandbox.googleapis.com\"\n",
        "else:\n",
        "    base_url = \"https://geminidataanalytics.googleapis.com\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "52oJ2poOfHCp"
      },
      "source": [
        "# Data Agent Services"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RuYxbQSmjSeD"
      },
      "outputs": [],
      "source": [
        "# @title Create Data Agent\n",
        "\n",
        "datasource_references = bigquery_data_sources\n",
        "# datasource_references = looker_data_sources\n",
        "# datasource_references = looker_studio_data_sources\n",
        "\n",
        "# fmt: off\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents\"\n",
        "\n",
        "data_agent_payload = {\n",
        "    \"name\": f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\",  # Optional\n",
        "    \"description\": \"This is the description of data_agent.\",  # Optional\n",
        "    \"data_analytics_agent\": {\n",
        "        \"published_context\": {\n",
        "            \"datasource_references\": datasource_references,\n",
        "            \"system_instruction\": system_instruction,\n",
        "            \"example_queries\": example_queries,  # Optional - Only supported for BigQuery\n",
        "            \"options\": {\n",
        "                \"analysis\": {\n",
        "                    # Optional - if wanting to use advanced analysis with python.\n",
        "                    # Default is False.\n",
        "                    \"python\": {\"enabled\": False}\n",
        "                }\n",
        "            },\n",
        "        }\n",
        "    },\n",
        "}\n",
        "\n",
        "params = {\"data_agent_id\": data_agent_id}  # Optional\n",
        "\n",
        "data_agent_response = requests.post(\n",
        "    data_agent_url, params=params, json=data_agent_payload, headers=headers\n",
        ")\n",
        "\n",
        "if data_agent_response.status_code == 200:\n",
        "    print(\"Data Agent created successfully!\")\n",
        "    print(json.dumps(data_agent_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error creating Data Agent: {data_agent_response.status_code}\")\n",
        "    print(data_agent_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gKeVQqsGkyqT"
      },
      "outputs": [],
      "source": [
        "# @title Get Data Agent\n",
        "\n",
        "# fmt: off\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n",
        "\n",
        "data_agent_response = requests.get(data_agent_url, headers=headers)\n",
        "\n",
        "if data_agent_response.status_code == 200:\n",
        "    print(\"Fetched Data Agent successfully!\")\n",
        "    print(json.dumps(data_agent_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error: {data_agent_response.status_code}\")\n",
        "    print(data_agent_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VP2oNORjo4bc"
      },
      "outputs": [],
      "source": [
        "# @title List Data Agents\n",
        "\n",
        "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents\"\n",
        "\n",
        "data_agent_response = requests.get(data_agent_url, headers=headers)\n",
        "\n",
        "if data_agent_response.status_code == 200:\n",
        "    print(\"Data Agents Listed successfully!\")\n",
        "    print(json.dumps(data_agent_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error Listing Data Agents: {data_agent_response.status_code}\")\n",
        "    print(data_agent_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "be0ebe371ad7"
      },
      "outputs": [],
      "source": [
        "# @title List Accessible Data Agents\n",
        "\n",
        "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents:listAccessible\"\n",
        "\n",
        "# Choose the desired creator filter\n",
        "# fmt: off\n",
        "creator_filter = \"NONE\"  # @param [\"NONE\", \"CREATOR_ONLY\", \"NOT_CREATOR_ONLY\"]\n",
        "# fmt: on\n",
        "\n",
        "# Add the creator_filter as a query parameter\n",
        "params = {\"creator_filter\": creator_filter}\n",
        "\n",
        "data_agent_response = requests.get(data_agent_url, headers=headers, params=params)\n",
        "\n",
        "if data_agent_response.status_code == 200:\n",
        "    print(\"Data Agents Listed successfully!\")\n",
        "    print(json.dumps(data_agent_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error Listing Data Agents: {data_agent_response.status_code}\")\n",
        "    print(data_agent_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-S7AIL1OeXnu"
      },
      "outputs": [],
      "source": [
        "# @title Update Data Agent\n",
        "\n",
        "# fmt: off\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n",
        "\n",
        "payload = {\n",
        "    \"description\": \"This is the latest description of data_agent.\",\n",
        "    \"data_analytics_agent\": {\n",
        "        \"published_context\": {\n",
        "            \"datasource_references\": bigquery_data_sources,\n",
        "            \"system_instruction\": system_instruction,\n",
        "        }\n",
        "    },\n",
        "}\n",
        "\n",
        "fields = [\"description\", \"data_analytics_agent\"]\n",
        "params = {\"updateMask\": \",\".join(fields)}\n",
        "\n",
        "data_agent_response = requests.patch(\n",
        "    data_agent_url, headers=headers, params=params, json=payload\n",
        ")\n",
        "\n",
        "if data_agent_response.status_code == 200:\n",
        "    print(\"Data Agent updated successfully!\")\n",
        "    print(json.dumps(data_agent_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error Updating Data Agent: {data_agent_response.status_code}\")\n",
        "    print(data_agent_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TjZfg2lVF8Jl"
      },
      "outputs": [],
      "source": [
        "# @title [Agent Sharing] Set IAM Policy for Data Agent\n",
        "\n",
        "\n",
        "\"\"\" PLEASE NOTE THIS API CALLS OVERRIDES EXISTING PERMISSION FOR THE RESOURCE.\n",
        "For preserving existing policy in practise call Get IAM policy to fetch existing policy and pass it along with additional changes\n",
        "in the call to Set IAM Policy\n",
        "\"\"\"\n",
        "# fmt: off\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "users = \"abc@google.com, wxyz@google.com\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}:setIamPolicy\"\n",
        "\n",
        "\n",
        "# Request body\n",
        "payload = {\n",
        "    \"policy\": {\n",
        "        \"bindings\": [\n",
        "            {\n",
        "                \"role\": \"roles/geminidataanalytics.dataAgentEditor\",\n",
        "                \"members\": [f\"user:{i.strip()}\" for i in users.split(\",\")],\n",
        "            }\n",
        "        ]\n",
        "    }\n",
        "}\n",
        "\n",
        "data_agent_response = requests.post(data_agent_url, headers=headers, json=payload)\n",
        "\n",
        "if data_agent_response.status_code == 200:\n",
        "    print(\"IAM Policy set successfully!\")\n",
        "    print(json.dumps(data_agent_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error setting IAM policy: {data_agent_response.status_code}\")\n",
        "    print(data_agent_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kHt_lhjh3LtS"
      },
      "outputs": [],
      "source": [
        "# @title [Agent Sharing] Get IAM Policy for Data Agent\n",
        "\n",
        "# fmt: off\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}:getIamPolicy\"\n",
        "\n",
        "# Request body\n",
        "payload = {\n",
        "    \"resource\": f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n",
        "}\n",
        "\n",
        "data_agent_response = requests.post(data_agent_url, headers=headers, json=payload)\n",
        "\n",
        "if data_agent_response.status_code == 200:\n",
        "    print(\"IAM Policy fetched successfully!\")\n",
        "    print(json.dumps(data_agent_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error fetching IAM policy: {data_agent_response.status_code}\")\n",
        "    print(data_agent_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "y-wlEDnPrOvO"
      },
      "outputs": [],
      "source": [
        "# @title [Soft Delete] Delete Data Agent\n",
        "\n",
        "# fmt: off\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n",
        "\n",
        "data_agent_response = requests.delete(data_agent_url, headers=headers)\n",
        "\n",
        "if data_agent_response.status_code == 200:\n",
        "    print(\"Data Agent deleted successfully!\")\n",
        "    print(json.dumps(data_agent_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error Deleting Data Agent: {data_agent_response.status_code}\")\n",
        "    print(data_agent_response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KHQUswQMHkcU"
      },
      "source": [
        "# Data Chat Service"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "c9aCGQAYCMx1"
      },
      "outputs": [],
      "source": [
        "# @title Create Conversation\n",
        "\n",
        "# fmt: off\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "conversation_id = \"conversation_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations\"\n",
        "\n",
        "conversation_payload = {\n",
        "    \"agents\": [\n",
        "        f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n",
        "    ],\n",
        "    \"name\": f\"projects/{billing_project}/locations/{location}/conversations/{conversation_id}\",\n",
        "}\n",
        "params = {\"conversation_id\": conversation_id}\n",
        "\n",
        "conversation_response = requests.post(\n",
        "    conversation_url, headers=headers, params=params, json=conversation_payload\n",
        ")\n",
        "\n",
        "if conversation_response.status_code == 200:\n",
        "    print(\"Conversation created successfully!\")\n",
        "    print(json.dumps(conversation_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error creating Conversation: {conversation_response.status_code}\")\n",
        "    print(conversation_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RxxC3Er3HpIR"
      },
      "outputs": [],
      "source": [
        "# @title Get Conversation\n",
        "\n",
        "# fmt: off\n",
        "conversation_id = \"conversation_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}\"\n",
        "\n",
        "conversation_response = requests.get(conversation_url, headers=headers)\n",
        "\n",
        "# Handle the response\n",
        "if conversation_response.status_code == 200:\n",
        "    print(\"Conversation fetched successfully!\")\n",
        "    print(json.dumps(conversation_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error while fetching conversation: {conversation_response.status_code}\")\n",
        "    print(conversation_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3-3otk_aL6Wi"
      },
      "outputs": [],
      "source": [
        "# @title List Conversation\n",
        "\n",
        "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations\"\n",
        "\n",
        "\n",
        "conversation_response = requests.get(conversation_url, headers=headers)\n",
        "\n",
        "# Handle the response\n",
        "if conversation_response.status_code == 200:\n",
        "    print(\"Conversation fetched successfully!\")\n",
        "    print(json.dumps(conversation_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error while fetching conversation: {conversation_response.status_code}\")\n",
        "    print(conversation_response.text)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Delete Conversation\n",
        "\n",
        "# fmt: off\n",
        "conversation_id = \"conversation_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}\"\n",
        "\n",
        "conversation_response = requests.delete(conversation_url, headers=headers)\n",
        "\n",
        "# Handle the response\n",
        "if conversation_response.status_code == 200:\n",
        "    print(\"Conversation deleted successfully!\")\n",
        "    print(json.dumps(conversation_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error while deleting conversation: {conversation_response.status_code}\")\n",
        "    print(conversation_response.text)"
      ],
      "metadata": {
        "id": "rnIQ5blAUIuV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1koXYj2eD8jW"
      },
      "outputs": [],
      "source": [
        "# @title Streaming Chat Messages\n",
        "\n",
        "\n",
        "def is_json(str):\n",
        "    try:\n",
        "        json_object = json_lib.loads(str)\n",
        "    except ValueError:\n",
        "        return False\n",
        "    return True\n",
        "\n",
        "def handle_text_response(resp):\n",
        "    parts = resp[\"parts\"]\n",
        "    full_text = \"\".join(parts)\n",
        "    if \"\\n\" not in full_text and len(full_text) > 80:\n",
        "        wrapped_text = textwrap.fill(full_text, width=80)\n",
        "        print(wrapped_text)\n",
        "    else:\n",
        "        print(full_text)\n",
        "\n",
        "def get_property(data, field_name, default=\"\"):\n",
        "    return data[field_name] if field_name in data else default\n",
        "\n",
        "def display_schema(data):\n",
        "    fields = data[\"fields\"]\n",
        "    df = pd.DataFrame(\n",
        "        {\n",
        "            \"Column\": map(lambda field: get_property(field, \"name\"), fields),\n",
        "            \"Type\": map(lambda field: get_property(field, \"type\"), fields),\n",
        "            \"Description\": map(\n",
        "                lambda field: get_property(field, \"description\", \"-\"), fields\n",
        "            ),\n",
        "            \"Mode\": map(lambda field: get_property(field, \"mode\"), fields),\n",
        "        }\n",
        "    )\n",
        "    display(df)\n",
        "\n",
        "def display_section_title(text):\n",
        "    display(HTML(f\"<h2>{text}</h2>\"))\n",
        "\n",
        "def format_bq_table_ref(table_ref):\n",
        "    return \"{}.{}.{}\".format(\n",
        "        table_ref[\"projectId\"], table_ref[\"datasetId\"], table_ref[\"tableId\"]\n",
        "    )\n",
        "\n",
        "def format_looker_table_ref(table_ref):\n",
        "    return \"lookmlModel: {}, explore: {}\".format(\n",
        "        table_ref[\"lookmlModel\"], table_ref[\"explore\"]\n",
        "    )\n",
        "\n",
        "def display_datasource(datasource):\n",
        "    source_name = \"\"\n",
        "\n",
        "    if \"studioDatasourceId\" in datasource:\n",
        "        source_name = datasource[\"studioDatasourceId\"]\n",
        "    elif \"lookerExploreReference\" in datasource:\n",
        "        source_name = format_looker_table_ref(datasource[\"lookerExploreReference\"])\n",
        "    else:\n",
        "        source_name = format_bq_table_ref(datasource[\"bigqueryTableReference\"])\n",
        "\n",
        "    print(source_name)\n",
        "    display_schema(datasource[\"schema\"])\n",
        "\n",
        "def handle_schema_response(resp):\n",
        "    if \"query\" in resp:\n",
        "        print(resp[\"query\"][\"question\"])\n",
        "    elif \"result\" in resp:\n",
        "        display_section_title(\"Schema resolved\")\n",
        "        print(\"Data sources:\")\n",
        "        for datasource in resp[\"result\"][\"datasources\"]:\n",
        "            display_datasource(datasource)\n",
        "\n",
        "def handle_data_response(resp):\n",
        "    if \"query\" in resp:\n",
        "        query = resp[\"query\"]\n",
        "        display_section_title(\"Retrieval query\")\n",
        "        print(\"Query name: {}\".format(query[\"name\"]))\n",
        "        if \"question\" in query:\n",
        "          print(\"Question: {}\".format(query[\"question\"]))\n",
        "        if \"datasources\" in query:\n",
        "          print(\"Data sources:\")\n",
        "          for datasource in query[\"datasources\"]:\n",
        "              display_datasource(datasource)\n",
        "    elif \"generatedSql\" in resp:\n",
        "        display_section_title(\"SQL generated\")\n",
        "        print(resp[\"generatedSql\"])\n",
        "    elif \"result\" in resp:\n",
        "        display_section_title(\"Data retrieved\")\n",
        "\n",
        "        fields = map(\n",
        "            lambda field: get_property(field, \"name\"),\n",
        "            resp[\"result\"][\"schema\"][\"fields\"],\n",
        "        )\n",
        "        dict = {}\n",
        "\n",
        "        for field in fields:\n",
        "            dict[field] = [get_property(el, field) for el in resp[\"result\"][\"data\"]]\n",
        "\n",
        "        display(pd.DataFrame(dict))\n",
        "\n",
        "def handle_chart_response(resp):\n",
        "    if \"query\" in resp:\n",
        "        print(resp[\"query\"][\"instructions\"])\n",
        "    elif \"result\" in resp:\n",
        "        vegaConfig = resp[\"result\"][\"vegaConfig\"]\n",
        "        alt.Chart.from_json(json_lib.dumps(vegaConfig)).display()\n",
        "\n",
        "def handle_error(resp):\n",
        "    display_section_title(\"Error\")\n",
        "    print(\"Code: {}\".format(resp[\"code\"]))\n",
        "    print(\"Message: {}\".format(resp[\"message\"]))\n",
        "\n",
        "def get_stream(url, json):\n",
        "    s = requests.Session()\n",
        "\n",
        "    acc = \"\"\n",
        "\n",
        "    with s.post(url, json=json, headers=headers, stream=True) as resp:\n",
        "        for line in resp.iter_lines():\n",
        "            if not line:\n",
        "                continue\n",
        "\n",
        "            decoded_line = str(line, encoding=\"utf-8\")\n",
        "\n",
        "            if decoded_line == \"[{\":\n",
        "                acc = \"{\"\n",
        "            elif decoded_line == \"}]\":\n",
        "                acc += \"}\"\n",
        "            elif decoded_line == \",\":\n",
        "                continue\n",
        "            else:\n",
        "                acc += decoded_line\n",
        "\n",
        "            if not is_json(acc):\n",
        "                continue\n",
        "\n",
        "            data_json = json_lib.loads(acc)\n",
        "\n",
        "            if \"systemMessage\" not in data_json:\n",
        "                if \"error\" in data_json:\n",
        "                    handle_error(data_json[\"error\"])\n",
        "                continue\n",
        "\n",
        "            if \"text\" in data_json[\"systemMessage\"]:\n",
        "                handle_text_response(data_json[\"systemMessage\"][\"text\"])\n",
        "            elif \"schema\" in data_json[\"systemMessage\"]:\n",
        "                handle_schema_response(data_json[\"systemMessage\"][\"schema\"])\n",
        "            elif \"data\" in data_json[\"systemMessage\"]:\n",
        "                handle_data_response(data_json[\"systemMessage\"][\"data\"])\n",
        "            elif \"chart\" in data_json[\"systemMessage\"]:\n",
        "                handle_chart_response(data_json[\"systemMessage\"][\"chart\"])\n",
        "            else:\n",
        "                colored_json = highlight(\n",
        "                    acc, lexers.JsonLexer(), formatters.TerminalFormatter()\n",
        "                )\n",
        "                print(colored_json)\n",
        "            print(\"\\n\")\n",
        "            acc = \"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "mQ1G8_UnKQjT"
      },
      "outputs": [],
      "source": [
        "# @title [Stateful] Chat using Conversation\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "conversation_id = \"conversation_1\"  # @param {type:\"string\"}\n",
        "# fmt: off\n",
        "question = \"Make a bar graph for the top 5 states by the total number of airports\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "chat_url = (\n",
        "    f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}:chat\"\n",
        ")\n",
        "\n",
        "# Construct the payload\n",
        "chat_payload = {\n",
        "    \"parent\": f\"projects/{billing_project}/locations/global\",\n",
        "    \"messages\": [{\"userMessage\": {\"text\": question}}],\n",
        "    \"conversation_reference\": {\n",
        "        \"conversation\": f\"projects/{billing_project}/locations/{location}/conversations/{conversation_id}\",\n",
        "        \"data_agent_context\": {\n",
        "            \"data_agent\": f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\",\n",
        "            # \"credentials\": looker_credentials #Uncomment this if your dataAgent contains looker as dataSource\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "# Call get_stream method to stream the response\n",
        "get_stream(chat_url, chat_payload)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "p0vGFMmzI8BH"
      },
      "outputs": [],
      "source": [
        "# @title [Stateless] Chat using Data Agent\n",
        "\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "# fmt: off\n",
        "question = \"Make a bar graph for the top 5 states by the total number of airports\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "chat_url = (\n",
        "    f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}:chat\"\n",
        ")\n",
        "\n",
        "# Construct the payload\n",
        "chat_payload = {\n",
        "    \"parent\": f\"projects/{billing_project}/locations/global\",\n",
        "    \"messages\": [{\"userMessage\": {\"text\": question}}],\n",
        "    \"data_agent_context\": {\n",
        "        \"data_agent\": f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\",\n",
        "        # \"credentials\": looker_credentials #Uncomment this if your dataAgent contains looker as dataSource\n",
        "    },\n",
        "}\n",
        "\n",
        "# Call get_stream method to stream the response\n",
        "get_stream(chat_url, chat_payload)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-S5ccMveLG3f"
      },
      "outputs": [],
      "source": [
        "# @title [Stateless] Chat using Inline Context\n",
        "\n",
        "chat_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/global:chat\"\n",
        "\n",
        "# fmt: off\n",
        "question = \"Show the schema of datasource\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "# Construct the payload\n",
        "chat_payload = {\n",
        "    \"parent\": f\"projects/{billing_project}/locations/global\",\n",
        "    \"messages\": [{\"userMessage\": {\"text\": question}}],\n",
        "    \"inline_context\": {\n",
        "        \"datasource_references\": bigquery_data_sources,\n",
        "        \"example_queries\": example_queries,  # Optional - Only supported for BigQuery\n",
        "        \"options\": {\n",
        "            \"analysis\": {\n",
        "                # Optional - if wanting to use advanced analysis with python.\n",
        "                # Default is False.\n",
        "                \"python\": {\"enabled\": False}\n",
        "            }\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "# Call get_stream method to stream the response\n",
        "get_stream(chat_url, chat_payload)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Ioj6gG8wMedz"
      },
      "outputs": [],
      "source": [
        "# @title List Messages\n",
        "\n",
        "# fmt: off\n",
        "conversation_id = \"conversation_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}/messages\"\n",
        "\n",
        "conversation_response = requests.get(conversation_url, headers=headers)\n",
        "\n",
        "# Handle the response\n",
        "if conversation_response.status_code == 200:\n",
        "    print(\"Conversation fetched successfully!\")\n",
        "    print(json.dumps(conversation_response.json(), indent=2))\n",
        "else:\n",
        "    print(f\"Error while fetching conversation: {conversation_response.status_code}\")\n",
        "    print(conversation_response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wdhjfwM3tF4G"
      },
      "source": [
        "## Multi-turn Conversation with Data Agent or Inline Context"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Lr3vhcG5tNjm"
      },
      "outputs": [],
      "source": [
        "# @title get_stream function for multi-turn conversation\n",
        "\n",
        "# NOTE: This methods is same as get_stream() method present in Streaming Chat Messages section.\n",
        "# The only difference is that - here we are storing the response in an array \"conversation_messages\" to save the conversation.\n",
        "\n",
        "\n",
        "def get_stream_multi_turn(url, json, conversation_messages):\n",
        "    s = requests.Session()\n",
        "\n",
        "    acc = \"\"\n",
        "\n",
        "    with s.post(url, json=json, headers=headers, stream=True) as resp:\n",
        "        for line in resp.iter_lines():\n",
        "            if not line:\n",
        "                continue\n",
        "\n",
        "            decoded_line = str(line, encoding=\"utf-8\")\n",
        "\n",
        "            if decoded_line == \"[{\":\n",
        "                acc = \"{\"\n",
        "            elif decoded_line == \"}]\":\n",
        "                acc += \"}\"\n",
        "            elif decoded_line == \",\":\n",
        "                continue\n",
        "            else:\n",
        "                acc += decoded_line\n",
        "\n",
        "            if not is_json(acc):\n",
        "                continue\n",
        "\n",
        "            data_json = json_lib.loads(acc)\n",
        "            # Store the response to be used in next iteration.\n",
        "            conversation_messages.append(data_json)\n",
        "\n",
        "            if \"systemMessage\" not in data_json:\n",
        "                if \"error\" in data_json:\n",
        "                    handle_error(data_json[\"error\"])\n",
        "                continue\n",
        "\n",
        "            if \"text\" in data_json[\"systemMessage\"]:\n",
        "                handle_text_response(data_json[\"systemMessage\"][\"text\"])\n",
        "            elif \"schema\" in data_json[\"systemMessage\"]:\n",
        "                handle_schema_response(data_json[\"systemMessage\"][\"schema\"])\n",
        "            elif \"data\" in data_json[\"systemMessage\"]:\n",
        "                handle_data_response(data_json[\"systemMessage\"][\"data\"])\n",
        "            elif \"chart\" in data_json[\"systemMessage\"]:\n",
        "                handle_chart_response(data_json[\"systemMessage\"][\"chart\"])\n",
        "            else:\n",
        "                colored_json = highlight(\n",
        "                    acc, lexers.JsonLexer(), formatters.TerminalFormatter()\n",
        "                )\n",
        "                print(colored_json)\n",
        "            print(\"\\n\")\n",
        "            acc = \"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "In8Gtf_Yte_i"
      },
      "outputs": [],
      "source": [
        "# @title Multi-turn Conversation helper function\n",
        "chat_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/global:chat\"\n",
        "\n",
        "# Re-used across requests to track previous turns\n",
        "conversation_messages = []\n",
        "\n",
        "# fmt: off\n",
        "data_agent_id = \"data_agent_1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "# Helper function for calling the API\n",
        "\n",
        "\n",
        "def multi_turn_conversation(msg):\n",
        "    userMessage = {\"userMessage\": {\"text\": msg}}\n",
        "\n",
        "    # Send multi-turn request by including previous turns, plus new message\n",
        "    conversation_messages.append(userMessage)\n",
        "\n",
        "    # Construct the payload\n",
        "    chat_payload = {\n",
        "        \"parent\": f\"projects/{billing_project}/locations/global\",\n",
        "        \"messages\": conversation_messages,\n",
        "        \"data_agent_context\": {\n",
        "            \"data_agent\": f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\",\n",
        "            # \"credentials\": looker_credentials # Uncomment this if your dataAgent contains looker as dataSource\n",
        "        },\n",
        "        # uncomment the lines below when using chat with inline_context\n",
        "        # \"inline_context\": {\n",
        "        #     \"datasource_references\": bigquery_data_sources,\n",
        "        # }\n",
        "    }\n",
        "\n",
        "    # Call get_stream_multi_turn method to stream the response\n",
        "    get_stream_multi_turn(chat_url, chat_payload, conversation_messages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M_8Lvo68thp4"
      },
      "outputs": [],
      "source": [
        "# Send first-turn request\n",
        "multi_turn_conversation(\"List of the top 5 states by the total number of airports\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RFSy-DNTtjf7"
      },
      "outputs": [],
      "source": [
        "# Send follow-up-turn request\n",
        "multi_turn_conversation(\"Can you please make it a pie chart?\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "intro_gemini_data_analytics_http.ipynb",
      "toc_visible": true,
      "provenance": [],
      "last_runtime": {
        "build_target": "//cloud/data_analytics/anarres/colab/cortado:base_colab",
        "kind": "private"
      }
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
